import os
import re
import subprocess
import csv_to_docx
from docx import Document
from docx.oxml.ns import qn
from PIL import Image
import io


def docx_to_markdown(docx_path, output_md_path, media_dir="media"):
    os.makedirs(media_dir, exist_ok=True)
    tables_dir = os.path.join(media_dir, "tables")
    os.makedirs(tables_dir, exist_ok=True)

    doc = Document(docx_path)
    md_content = []
    image_count = 0
    table_count = 0

    for element in doc.element.body:
        if element.tag.endswith('p'):
            for para in doc.paragraphs:
                if para._element is element:
                    # 处理段落文本
                    if para.text.strip():
                        md_content.append(process_paragraph(para))

                    # 处理图片
                    images_md = process_images_in_paragraph(para, media_dir, image_count)
                    if images_md:
                        md_content.append(images_md)
                        image_count += images_md.count("![图片描述]")
                    break

        elif element.tag.endswith('tbl'):
            for table in doc.tables:
                if table._element is element:
                    table_filename = f"table_{table_count}.csv"
                    table_path = os.path.join(tables_dir, table_filename)
                    table_md = process_table(table, table_path, output_md_path)
                    if table_md:
                        md_content.append(table_md)
                        table_count += 1
                    break

    with open(output_md_path, 'w', encoding='utf-8') as md_file:
        md_file.write("\n\n".join([c for c in md_content if c]))

    print(f"Markdown文件已保存至: {output_md_path}")
    print(f"图片已保存至: {media_dir}")
    print(f"表格已保存至: {tables_dir}")


def process_paragraph(para):
    text = para.text.strip()
    if not text:
        return ""

    # 标题样式映射
    level_map = {
        'Heading 1': 1, 'Heading 2': 2, 'Heading 3': 3,
        'Heading 4': 4, 'Heading 5': 5, 'Heading 6': 6,
        'T-2': 2, 'T-3': 3  # 添加自定义样式映射
    }

    if para.style.name in level_map:
        level = level_map[para.style.name]

        # 提取并处理锚点
        anchor_match = re.search(r'\{#([\w-]+)\}', text)
        anchor_id = anchor_match.group(1) if anchor_match else ""
        text = re.sub(r'\s*\{#[\w-]+\}\s*', '', text)  # 移除锚点标记

        if anchor_id:
            return f"{'#' * level} {text} {{#{anchor_id}}}"
        else:
            return f"{'#' * level} {text}"

    # 处理列表
    if para.style.name == 'List Paragraph' or para.style.name.startswith('List'):
        return f"- {text}"

    # 保留原始换行符
    text = text.replace('\n', '  \n')
    return text


def process_images_in_paragraph(para, media_dir, img_index):
    md_images = []
    for run in para.runs:
        drawing = run._element.find('.//' + qn('w:drawing'))
        if drawing is not None:
            blip = drawing.find('.//' + qn('a:blip'))
            if blip is not None:
                try:
                    embed_attr = qn('r:embed')
                    image_id = blip.get(embed_attr)
                    if image_id:
                        image_part = run.part.related_parts[image_id]
                        image_bytes = image_part.blob
                        try:
                            img = Image.open(io.BytesIO(image_bytes))
                            image_ext = img.format.lower()
                        except:
                            image_ext = "png"
                        image_name = f"image{img_index}.{image_ext}"
                        image_path = os.path.join(media_dir, image_name)
                        with open(image_path, 'wb') as img_file:
                            img_file.write(image_bytes)
                        md_images.append(f"![图片描述]({image_path})")
                        img_index += 1
                except Exception as e:
                    print(f"图片处理错误: {str(e)}")
                    md_images.append("[图片处理错误]")
    return "\n".join(md_images)


def process_table(table, csv_path, md_path):
    with open(csv_path, 'w', encoding='utf-8', newline='') as csv_file:
        writer = csv.writer(csv_file)
        for row in table.rows:
            row_data = [cell.text.strip().replace('\n', ' ') for cell in row.cells]
            writer.writerow(row_data)
    md_dir = os.path.dirname(md_path)
    rel_csv_path = os.path.relpath(csv_path, md_dir)
    return f"[点击查看表格内容]({rel_csv_path})"


def enhance_markdown_with_pandoc(input_md, output_md):
    try:
        subprocess.run([
            'pandoc', '-s', input_md,
            '-t', 'markdown_strict+pipe_tables',
            '-o', output_md
        ], check=True)
        print(f"使用pandoc优化后的Markdown已保存至: {output_md}")
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("未找到pandoc，使用基本转换结果")


def docx_to_md(path):
    # 提取文件名（不带扩展名）
    filename = os.path.splitext(os.path.basename(path))[0]
    # 要拼接的目标文件夹
    output_md = os.path.join(r"D:\hbyt\AI智能投标\markdown_output", filename, f"{filename}.md")
    enhanced_md = os.path.join(r"D:\hbyt\AI智能投标\markdown_output", f"{filename}_enhanced.md")
    media_folder = os.path.join(r"D:\hbyt\AI智能投标\markdown_output\media", filename)
    tables_folder = os.path.join(r"D:\hbyt\AI智能投标\markdown_output\tables", filename)
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_md), exist_ok=True)
    os.makedirs(media_folder, exist_ok=True)
    docx_to_markdown(docx_file, output_md, media_folder)
    enhance_markdown_with_pandoc(output_md, enhanced_md)
    return output_md

if __name__ == "__main__":
    docx_file = "D:\\hbyt\\AI智能投标\\能力组件-附件\\激励及绩效管理_v1.0_2201.docx"
    # docx_file = "D:\\hbyt\\AI智能投标\\能力组件-附件\\EHS保障体系_v1.0_2211.docx"
    output_md = "D:\\hbyt\\AI智能投标\\markdown_output\\markdown_output1.md"
    enhanced_md = "D:\\hbyt\\AI智能投标\\markdown_output\\enhanced.md"
    media_folder = "D:\\hbyt\\AI智能投标\\markdown_output\\media"
    os.makedirs(os.path.dirname(output_md), exist_ok=True)
    os.makedirs(media_folder, exist_ok=True)
    docx_to_markdown(docx_file, output_md, media_folder)
    enhance_markdown_with_pandoc(output_md, enhanced_md)